package core

import (
	"sync"

	"github.com/panjf2000/ants/v2"
)

type ProgressCallback func(taskID string, progress float64)

// TaskManager 管理所有 Task
type TaskManager struct {
	tasks    map[string]*Task
	mu       sync.RWMutex
	pool     *ants.Pool
	file     *StatusFileStorage
	cache    *Cache
	callback ProgressCallback
}

// NewTaskManager 创建 TaskManager
func NewTaskManager(poolSize int, output string) (*TaskManager, error) {
	pool, err := ants.NewPool(poolSize)
	if err != nil {
		return nil, err
	}

	return &TaskManager{
		tasks: make(map[string]*Task),
		pool:  pool,
		file:  NewStatusFileStorage(output),
		cache: NewCache(),
	}, nil
}

// AddTask 添加 Task
func (m *TaskManager) AddTask(t *Task) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.tasks[t.ID] = t
}

// GetTask 获取 Task
func (m *TaskManager) GetTask(taskID string) (*Task, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	t, ok := m.tasks[taskID]
	return t, ok
}

// RemoveTask 删除 Task
func (m *TaskManager) RemoveTask(taskID string) {
	m.mu.Lock()
	defer m.mu.Unlock()
	delete(m.tasks, taskID)
}

func (m *TaskManager) SetProgressCallback(cb ProgressCallback) {
	m.callback = cb
}

// RunTask 并发执行 Task 的所有 Job
func (m *TaskManager) RunTask(t *Task) {
	t.mu.Lock()
	t.Status = TaskRunning
	t.mu.Unlock()

	var wg sync.WaitGroup
	for _, job := range t.Jobs {
		j := job
		wg.Add(1)
		_ = m.pool.Submit(func() {
			defer wg.Done()
			j.Run()
			t.UpdateProgress()
			if m.callback != nil {
				m.callback(t.ID, t.Progress)
			}
		})
	}

	wg.Wait()
	t.UpdateProgress()
	if m.callback != nil {
		m.callback(t.ID, t.Progress)
	}
}

// CancelTask 取消 Task
func (m *TaskManager) CancelTask(taskID string) {
	t, ok := m.GetTask(taskID)
	if ok {
		t.CancelTask()
	}
}
